from managers.nodes.baseGraphNode import BaseGraphNode, Msg
from .sessionInfo import SessionInfo
from ..execute.baseExeHis import BaseExeHis
import requests
import json
from typing import Dict, Any, List

class KnowledgeGraphNode(BaseGraphNode):
    """
    百度 VDB 知识库查询节点
    支持：1. 动态查询参数输入  2. 返回结构化 JSON / 文本摘要
    """
    def __init__(self, id: str, data: dict, graphNodeManager, random_key: str):
        super().__init__(id, data, graphNodeManager, random_key)
        # === 前端配置项 ===
        self.api_url = data['data'].get('apiUrl', 'https://vdb.bj.baidubce.com/v1')
        self.collection = data['data']['collection']          # 必填：库名称
        self.top_k = data['data'].get('topK', 5)              # 返回条数
        self.threshold = data['data'].get('threshold', 0.7)   # 相似度阈值
        self.ak = data['data']['accessKey']                   # 必填：百度云 AK
        self.sk = data['data']['secretKey']                   # 必填：百度云 SK
        self.running = False

    # ---------- 百度 VDB 官方签名 ----------
    def _auth_headers(self, method: str, uri: str, params: str = '') -> Dict[str, str]:
        """简单版 VDB 签名，生产环境请用官方 SDK"""
        import hmac, hashlib, base64, time, uuid
        timestamp = str(int(time.time()))
        nonce = str(uuid.uuid4())
        sign_str = f'{method}\n{uri}\n{timestamp}\n{nonce}'
        signature = base64.b64encode(
            hmac.new(self.sk.encode(), sign_str.encode(), hashlib.sha256).digest()
        ).decode()
        return {
            'Content-Type': 'application/json',
            'x-bce-timestamp': timestamp,
            'x-bce-nonce': nonce,
            'Authorization': f'BCE {self.ak}:{signature}'
        }

    # ---------- 查询 VDB ----------
    def _query_vdb(self, query_text: str) -> List[Dict[str, Any]]:
        uri = f'/vdb/collections/{self.collection}/search'
        url = self.api_url.rstrip('/') + uri
        payload = {
            "vector": [],               # 如用文本向量，需先调 embedding 接口
            "text": query_text,         # VDB 支持直接文本相似度
            "topK": self.top_k,
            "score": self.threshold
        }
        headers = self._auth_headers('POST', uri)
        resp = requests.post(url, json=payload, headers=headers, timeout=10)
        resp.raise_for_status()
        return resp.json().get('results', [])

    # ---------- 节点主入口 ----------
    def execute(self, session_info: SessionInfo = None,
                exe_his: BaseExeHis = None, trigger=None, running_id=None):
        if session_info is None:
            return
        # 1. 取上游问题
        query_text = ''
        if 'query' in self.dict['in']:
            src = self.dict['in']['query']
            pre = self.graphNodeManager.nodesByIdDict[src['nodeId']]
            msg = pre.get_newest_output_in_session(session_info, src['handle'])
            if msg:
                query_text = msg.content

        if not query_text:
            results = []
        else:
            exe_his.update_content(f'VDB 查询: {query_text}')
            try:
                results = self._query_vdb(query_text)
            except Exception as e:
                results = [{'error': str(e)}]

        # 2. 输出
        with session_info.out_write_lock:
            if session_info.out.get(self.id) is None:
                session_info.out[self.id] = {}
            session_info.out[self.id]['results'] = [Msg(json.dumps(results), 'assistant', 'text')]
            session_info.out[self.id]['text'] = [Msg('\n'.join([r.get('text', '') for r in results]), 'assistant', 'text')]

        exe_his.done()
        self.flow_next(session_info, trigger=trigger, running_id=running_id, type='text', role='assistant')

    def stop(self):
        self.running = False

    def clear(self, session_info: SessionInfo):
        with session_info.out_write_lock:
            if session_info.out.get(self.id):
                session_info.out[self.id].clear()